import sys
import glob
import os
import time
import random
import argparse
import threading
import socket
import cv2
import numpy as np
try:
    sys.path.append(glob.glob('../carla/dist/carla-*%d.%d-%s.egg' % (
        sys.version_info.major,
        sys.version_info.minor,
        'win-amd64' if os.name == 'nt' else 'linux-x86_64'))[0])
except IndexError:
    pass
import carla


# 设置UDP客户端
UDP_IP = "192.168.0.2"  # 树莓派的IP地址
UDP_PORT = 5005         # 选择一个端口

# 创建保存图像的文件夹
if not os.path.exists('png'):
    os.makedirs('png')

def main():
    # 创建UDP套接字
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)

    argparser = argparse.ArgumentParser(
        description=__doc__)
    argparser.add_argument(
        '--host',
        metavar='H',
        default='127.0.0.1',
        help='IP of the host server (default: 127.0.0.1)')
    argparser.add_argument(
        '-p', '--port',
        metavar='P',
        default=2000,
        type=int,
        help='TCP port to listen to (default: 2000)')
    argparser.add_argument(
        '-s', '--speed',
        metavar='FACTOR',
        default=1.0,
        type=float,
        help='rate at which the weather changes (default: 1.0)')
    args = argparser.parse_args()

    # 连接到CARLA服务器
    client = carla.Client(args.host, args.port)
    client.set_timeout(10.0)
    world = client.get_world()

    # 获取可用的车辆模型
    blueprints = world.get_blueprint_library().filter('vehicle')
    vehicle_bp = [bp for bp in blueprints if 'tesla.model3' in bp.id][0]  # 选了个特斯拉model3

    # 获取所有可用的生成点（车道位置）
    spawn_points = world.get_map().get_spawn_points()
    
    # 随机选择一个生成点
    spawn_point = random.choice(spawn_points)

    # 在选定位置生成车辆
    vehicle = world.spawn_actor(vehicle_bp, spawn_point)

    # 创建摄像头蓝图
    camera_bp = world.get_blueprint_library().find('sensor.camera.rgb')
    camera_bp.set_attribute('image_size_x', '512')
    camera_bp.set_attribute('image_size_y', '256')
    camera_bp.set_attribute('fov', '120')

    # 假设挂载摄像头到车辆的前方
    #camera_transform = carla.Transform(carla.Location(x=1.5, z=3.0))  # 偏移量可以根据需要调整
    #carla.Rotation(pitch=-45)  # 向下倾斜 10 度（默认是0，负值表示向下）
    camera_transform = carla.Transform(
        carla.Location(x=1.6, z=1.7),  # 提高摄像头的高度 (z)
        carla.Rotation(pitch=-12)  # 向下倾斜 10 度（默认是0，负值表示向下）
    )
    # 在车辆上附加摄像头
    camera = world.spawn_actor(camera_bp, camera_transform, attach_to=vehicle)

    # 回调函数来处理图像数据
    def process_image(image):
        # 将原始图像数据转换为 OpenCV 格式
        image_array = np.frombuffer(image.raw_data, dtype=np.uint8)
        image_array = np.reshape(image_array, (image.height, image.width, 4))
        image_array = image_array[:, :, :3]  # 去掉透明度通道

        # 转换为BGR格式（OpenCV使用BGR）
        cv_image = cv2.cvtColor(image_array, cv2.COLOR_RGB2BGR)

        # 将图像编码为 JPEG 格式
        _, jpeg_image = cv2.imencode('.jpg', image_array)
        image_data = jpeg_image.tobytes()  # 将JPEG编码的图像转换为字节数组

        # 获取 JPEG 编码后的图像大小（单位：字节）
        image_size = len(image_data)
        print(f"发送的图像大小: {image_size} 字节")

        # 设置每个数据包的大小（以1440字节为例）
        MAX_PACKET_SIZE = 1440
        
        # 分割数据并发送
        for i in range(0, len(image_data), MAX_PACKET_SIZE):
            # 获取当前数据包
            packet = image_data[i:i + MAX_PACKET_SIZE]
            
            # 发送数据包
            sock.sendto(packet, (UDP_IP, UDP_PORT))
            print(f"Sent packet {i // MAX_PACKET_SIZE + 1}")

        # 保存图像为PNG格式
        timestamp = int(time.time())  # 使用时间戳作为文件名的一部分
        png_filename = f"png/carla_image_{timestamp}.png"
        cv2.imwrite(png_filename, cv_image)  # 保存图像为PNG格式
        print(f"图像已保存为 {png_filename}")

    # 设置摄像头回调
    camera.listen(lambda image: process_image(image))

    # 主循环
    try:
        while True:
            time.sleep(0.1)

    except KeyboardInterrupt:
        print("\n退出中...")
        
    finally:
        # 清理过程
        print("Destroying vehicle and camera...")
        camera.stop()  # 停止摄像头
        camera.destroy()  # 销毁摄像头
        vehicle.destroy()  # 销毁车辆
        cv2.destroyAllWindows()  # 关闭OpenCV窗口
        sock.close()
        print("Cleanup complete.")

if __name__ == '__main__':
    main()
